"""
@Author  : 吕申凯
@Time    : 2022/9/18 17:21
@File    : dropout_layer.py
@Function: 
"""
import torch
from torch import nn


class Dropout_net(nn.Module):
    """
        dropout网络
    """

    def __init__(self, neural_num, d_prob):
        """

        :param neural_num: 输入维度
        :param d_prob: dropout随机概率
        """
        super(Dropout_net, self).__init__()
        # 一层线性层，添加dropout正则
        self.linears = nn.Sequential(
            nn.Dropout(d_prob),  # 添加dropout正则，添加在网络层之前
            nn.Linear(neural_num, 1, bias=False),
            nn.ReLU(inplace=True)
        )

    def forward(self, input_X):
        """
            前项传播
        :param input_X:
        :return:
        """
        return self.linears(input_X)


if __name__ == '__main__':
    # 设置10000维数值为1的向量
    input_num = 10000
    x = torch.ones((input_num,), dtype=torch.float32)

    net = Dropout_net(input_num, d_prob=0.5)
    net.linears[1].weight.detach().fill_(1.)  # 权重设置为1

    # 训练
    '''
        10000个1经过一个权值全为1的神经网络，转化为一个数。
        测试环境下没有添加dropout正则化，最终结果期望为10000，
        训练环境添加0.5的概率使神经元失活，则训练集的的期望输出为5000。
        由于在训练环境下失活一些神经单元后，将下一层得到的值除上(1-失活概率)，
        这样最终得到的训练集的值和测试集在一个量级上，没有显示出0.5的差值。
    '''
    net.train()
    y = net(x)
    print("output in training mode", y)  # 10000左右

    # 测试
    net.eval()
    y = net(x)
    print("output in eval mode", y)  # 10000
